// Copyright 2014 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <memory>
#include <utility>
#include <vector>

#include "base/functional/bind.h"
#include "base/run_loop.h"
#include "base/task/single_thread_task_runner.h"
#include "base/test/task_environment.h"
#include "chromecast/media/api/test/mock_cma_backend.h"
#include "chromecast/media/base/decrypt_context_impl.h"
#include "chromecast/media/cdm/cast_cdm_context.h"
#include "chromecast/media/cma/pipeline/av_pipeline_client.h"
#include "chromecast/media/cma/pipeline/media_pipeline_impl.h"
#include "chromecast/media/cma/pipeline/video_pipeline_client.h"
#include "chromecast/media/cma/test/frame_generator_for_test.h"
#include "chromecast/media/cma/test/mock_frame_provider.h"
#include "chromecast/public/media/cast_decoder_buffer.h"
#include "media/base/audio_decoder_config.h"
#include "media/base/callback_registry.h"
#include "media/base/media_util.h"
#include "media/base/video_decoder_config.h"
#include "testing/gtest/include/gtest/gtest.h"

using testing::_;
using testing::AtLeast;
using testing::Invoke;
using testing::NiceMock;
using testing::Return;
using testing::SaveArg;

namespace {
// Total number of frames generated by CodedFrameProvider.
// The first frame has config, while the last one is EOS.
const int kNumFrames = 100;
const int kFrameSize = 512;
const int kFrameDurationUs = 40 * 1000;
const int kLastFrameTimestamp = (kNumFrames - 2) * kFrameDurationUs;
}  // namespace

namespace chromecast {
namespace media {

ACTION_P2(PushBuffer, delegate, buffer_pts) {
  if (arg0->end_of_stream()) {
    base::SingleThreadTaskRunner::GetCurrentDefault()->PostTask(
        FROM_HERE, base::BindOnce(&CmaBackend::Decoder::Delegate::OnEndOfStream,
                                  base::Unretained(*delegate)));
  } else {
    *buffer_pts = arg0->timestamp();
  }
  return CmaBackend::BufferStatus::kBufferSuccess;
}

class CastCdmContextForTest : public CastCdmContext {
 public:
  CastCdmContextForTest() : license_installed_(false) {}

  CastCdmContextForTest(const CastCdmContextForTest&) = delete;
  CastCdmContextForTest& operator=(const CastCdmContextForTest&) = delete;

  void SetLicenseInstalled() {
    license_installed_ = true;
    event_callbacks_.Notify(
        ::media::CdmContext::Event::kHasAdditionalUsableKey);
  }

  // CastCdmContext implementation:
  std::unique_ptr<::media::CallbackRegistration> RegisterEventCB(
      ::media::CdmContext::EventCB event_cb) override {
    return event_callbacks_.Register(std::move(event_cb));
  }

  std::unique_ptr<DecryptContextImpl> GetDecryptContext(
      const std::string& key_id,
      EncryptionScheme encryption_scheme) override {
    if (license_installed_) {
      return std::unique_ptr<DecryptContextImpl>(
          new DecryptContextImpl(KEY_SYSTEM_CLEAR_KEY));
    }
    return nullptr;
  }

  void SetKeyStatus(const std::string& key_id,
                    CastKeyStatus key_status,
                    uint32_t system_code) override {}

  void SetVideoResolution(int width, int height) override {}

 private:
  bool license_installed_;
  ::media::CallbackRegistry<::media::CdmContext::EventCB::RunType>
      event_callbacks_;
};

// Helper class for managing pipeline setup, teardown, feeding data, stop/start
// etc in a simple API for tests to use.
class PipelineHelper {
 public:
  enum Stream { STREAM_AUDIO, STREAM_VIDEO };

  PipelineHelper(bool audio, bool video, bool encrypted)
      : have_audio_(audio),
        have_video_(video),
        encrypted_(encrypted),
        pipeline_backend_(nullptr),
        audio_decoder_delegate_(nullptr),
        video_decoder_delegate_(nullptr) {}

  PipelineHelper(const PipelineHelper&) = delete;
  PipelineHelper& operator=(const PipelineHelper&) = delete;

  void Setup() {
    if (encrypted_) {
      cdm_context_.reset(new CastCdmContextForTest());
    }

    auto backend = std::make_unique<MockCmaBackend>();
    pipeline_backend_ = backend.get();
    ON_CALL(*pipeline_backend_, SetPlaybackRate(_)).WillByDefault(Return(true));
    ON_CALL(audio_decoder_, SetConfig(_)).WillByDefault(Return(true));
    ON_CALL(audio_decoder_, PushBuffer(_))
        .WillByDefault(PushBuffer(&audio_decoder_delegate_,
                                  &last_push_pts_[STREAM_AUDIO]));
    ON_CALL(video_decoder_, SetConfig(_)).WillByDefault(Return(true));
    ON_CALL(video_decoder_, PushBuffer(_))
        .WillByDefault(PushBuffer(&video_decoder_delegate_,
                                  &last_push_pts_[STREAM_VIDEO]));

    media_pipeline_ = std::make_unique<MediaPipelineImpl>();
    media_pipeline_->Initialize(kLoadTypeURL, std::move(backend),
                                /* is_buffering_enabled */ true);

    if (have_audio_) {
      ::media::AudioDecoderConfig audio_config(
          ::media::AudioCodec::kMP3, ::media::kSampleFormatS16,
          ::media::CHANNEL_LAYOUT_STEREO, 44100, ::media::EmptyExtraData(),
          ::media::EncryptionScheme::kUnencrypted);
      AvPipelineClient client;
      client.eos_cb = base::BindRepeating(&PipelineHelper::OnEos,
                                          base::Unretained(this), STREAM_AUDIO);
      EXPECT_CALL(*pipeline_backend_, CreateAudioDecoder())
          .Times(1)
          .WillOnce(Return(&audio_decoder_));
      EXPECT_CALL(audio_decoder_, SetDelegate(_))
          .Times(1)
          .WillOnce(SaveArg<0>(&audio_decoder_delegate_));
      ::media::PipelineStatus status = media_pipeline_->InitializeAudio(
          audio_config, std::move(client), CreateFrameProvider());
      ASSERT_EQ(::media::PIPELINE_OK, status);
    }
    if (have_video_) {
      std::vector<::media::VideoDecoderConfig> video_configs;
      video_configs.push_back(::media::VideoDecoderConfig(
          ::media::VideoCodec::kH264, ::media::H264PROFILE_MAIN,
          ::media::VideoDecoderConfig::AlphaMode::kIsOpaque,
          ::media::VideoColorSpace(), ::media::kNoTransformation,
          gfx::Size(640, 480), gfx::Rect(0, 0, 640, 480), gfx::Size(640, 480),
          ::media::EmptyExtraData(), ::media::EncryptionScheme()));
      VideoPipelineClient client;
      client.av_pipeline_client.eos_cb = base::BindRepeating(
          &PipelineHelper::OnEos, base::Unretained(this), STREAM_VIDEO);
      EXPECT_CALL(*pipeline_backend_, CreateVideoDecoder())
          .Times(1)
          .WillOnce(Return(&video_decoder_));
      EXPECT_CALL(video_decoder_, SetDelegate(_))
          .Times(1)
          .WillOnce(SaveArg<0>(&video_decoder_delegate_));
      ::media::PipelineStatus status = media_pipeline_->InitializeVideo(
          video_configs, std::move(client), CreateFrameProvider());
      ASSERT_EQ(::media::PIPELINE_OK, status);
    }
  }

  void SetPipelineStartExpectations() {
    // The pipeline will be paused first, for the initial data buffering. Then
    // it will be resumed, once enough data is buffered to start playback.
    // When starting media pipeline, GetCurrentPts will be called every
    // kTimeUpdateInterval(250ms).
    EXPECT_CALL(*pipeline_backend_, GetCurrentPts()).Times(AtLeast(1));
    EXPECT_CALL(*pipeline_backend_, Pause());
    EXPECT_CALL(*pipeline_backend_, SetPlaybackRate(1.0f));
    EXPECT_CALL(*pipeline_backend_, Resume());
  }

  // This is used for the Flush test case, where the pipeline start sequence is
  // interrupted by the Flush, and the initial buffering never completes.
  void SetPipelineStartFlushExpectations() {
    EXPECT_CALL(*pipeline_backend_, GetCurrentPts());
    EXPECT_CALL(*pipeline_backend_, Pause());
  }

  void Start(base::RepeatingClosure eos_cb) {
    eos_cb_ = std::move(eos_cb);
    eos_[STREAM_AUDIO] = !media_pipeline_->HasAudio();
    eos_[STREAM_VIDEO] = !media_pipeline_->HasVideo();
    last_push_pts_[STREAM_AUDIO] = std::numeric_limits<int64_t>::min();
    last_push_pts_[STREAM_VIDEO] = std::numeric_limits<int64_t>::min();
    int64_t start_pts = 0;

    EXPECT_CALL(*pipeline_backend_, Initialize())
        .Times(1)
        .WillOnce(Return(true));
    EXPECT_CALL(*pipeline_backend_, Start(start_pts))
        .Times(1)
        .WillOnce(Return(true));

    media_pipeline_->StartPlayingFrom(base::Milliseconds(start_pts));
    media_pipeline_->SetPlaybackRate(1.0f);
  }
  void SetCdm() { media_pipeline_->SetCdm(cdm_context_.get()); }
  void Flush(base::OnceClosure flush_cb) {
    EXPECT_CALL(*pipeline_backend_, Stop()).Times(1);
    media_pipeline_->Flush(std::move(flush_cb));
  }
  void Stop() {
    media_pipeline_.reset();
    base::RunLoop::QuitCurrentWhenIdleDeprecated();
  }
  void FlushThenStop() {
    base::OnceClosure stop_task =
        base::BindOnce(&PipelineHelper::Stop, base::Unretained(this));
    Flush(std::move(stop_task));
  }
  void SetCdmLicenseInstalled() { cdm_context_->SetLicenseInstalled(); }

  bool have_audio() const { return have_audio_; }
  bool have_video() const { return have_video_; }
  int64_t last_push_pts(Stream stream) const { return last_push_pts_[stream]; }

 private:
  std::unique_ptr<CodedFrameProvider> CreateFrameProvider() {
    std::vector<FrameGeneratorForTest::FrameSpec> frame_specs;
    frame_specs.resize(kNumFrames);
    for (size_t k = 0; k < frame_specs.size() - 1; k++) {
      frame_specs[k].has_config = (k == 0);
      frame_specs[k].timestamp = base::Microseconds(kFrameDurationUs) * k;
      frame_specs[k].size = kFrameSize;
      frame_specs[k].has_decrypt_config = encrypted_;
    }
    frame_specs.back().is_eos = true;

    std::unique_ptr<FrameGeneratorForTest> frame_generator(
        new FrameGeneratorForTest(frame_specs));
    bool provider_delayed_pattern[] = {false, true};
    std::unique_ptr<MockFrameProvider> frame_provider(new MockFrameProvider());
    frame_provider->Configure(
        std::vector<bool>(
            provider_delayed_pattern,
            provider_delayed_pattern + std::size(provider_delayed_pattern)),
        std::move(frame_generator));
    frame_provider->SetDelayFlush(true);
    return std::move(frame_provider);
  }

  void OnEos(Stream stream) {
    eos_[stream] = true;
    if (eos_[STREAM_AUDIO] && eos_[STREAM_VIDEO] && !eos_cb_.is_null())
      eos_cb_.Run();
  }

  bool have_audio_;
  bool have_video_;
  bool encrypted_;
  bool eos_[2];
  int64_t last_push_pts_[2];
  base::RepeatingClosure eos_cb_;
  std::unique_ptr<CastCdmContextForTest> cdm_context_;
  MockCmaBackend* pipeline_backend_;
  NiceMock<MockCmaBackend::AudioDecoder> audio_decoder_;
  NiceMock<MockCmaBackend::VideoDecoder> video_decoder_;
  CmaBackend::Decoder::Delegate* audio_decoder_delegate_;
  CmaBackend::Decoder::Delegate* video_decoder_delegate_;
  std::unique_ptr<MediaPipelineImpl> media_pipeline_;
};

using AudioVideoTuple = ::testing::tuple<bool, bool>;

class AudioVideoPipelineImplTest
    : public ::testing::TestWithParam<AudioVideoTuple> {
 public:
  AudioVideoPipelineImplTest() {}

  AudioVideoPipelineImplTest(const AudioVideoPipelineImplTest&) = delete;
  AudioVideoPipelineImplTest& operator=(const AudioVideoPipelineImplTest&) =
      delete;

 protected:
  void SetUp() override {
    pipeline_helper_.reset(new PipelineHelper(
        ::testing::get<0>(GetParam()), ::testing::get<1>(GetParam()), false));
    pipeline_helper_->Setup();
  }

  base::test::TaskEnvironment task_environment_;
  std::unique_ptr<PipelineHelper> pipeline_helper_;
};

static void VerifyPlay(PipelineHelper* pipeline_helper) {
  // The decoders must have received the last frame.
  if (pipeline_helper->have_audio())
    EXPECT_EQ(kLastFrameTimestamp,
              pipeline_helper->last_push_pts(PipelineHelper::STREAM_AUDIO));
  if (pipeline_helper->have_video())
    EXPECT_EQ(kLastFrameTimestamp,
              pipeline_helper->last_push_pts(PipelineHelper::STREAM_VIDEO));

  pipeline_helper->Stop();
}

TEST_P(AudioVideoPipelineImplTest, Play) {
  base::RepeatingClosure verify_task = base::BindRepeating(
      &VerifyPlay, base::Unretained(pipeline_helper_.get()));
  pipeline_helper_->SetPipelineStartExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(verify_task)));
  base::RunLoop().Run();
}

static void VerifyFlush(PipelineHelper* pipeline_helper) {
  // The decoders must not have received any frame.
  if (pipeline_helper->have_audio())
    EXPECT_LT(pipeline_helper->last_push_pts(PipelineHelper::STREAM_AUDIO), 0);
  if (pipeline_helper->have_video())
    EXPECT_LT(pipeline_helper->last_push_pts(PipelineHelper::STREAM_VIDEO), 0);

  pipeline_helper->Stop();
}

static void VerifyNotReached() {
  EXPECT_TRUE(false);
}

TEST_P(AudioVideoPipelineImplTest, Flush) {
  base::OnceClosure verify_task =
      base::BindOnce(&VerifyFlush, base::Unretained(pipeline_helper_.get()));
  pipeline_helper_->SetPipelineStartFlushExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                base::BindRepeating(&VerifyNotReached)));
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Flush,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(verify_task)));

  base::RunLoop().Run();
}

TEST_P(AudioVideoPipelineImplTest, FullCycle) {
  base::RepeatingClosure eos_cb = base::BindRepeating(
      &PipelineHelper::FlushThenStop, base::Unretained(pipeline_helper_.get()));

  pipeline_helper_->SetPipelineStartExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(eos_cb)));
  base::RunLoop().Run();
}

// Test all three types of pipeline: audio-only, video-only, audio-video.
INSTANTIATE_TEST_SUITE_P(
    MediaPipelineImplTests,
    AudioVideoPipelineImplTest,
    ::testing::Values(AudioVideoTuple(true, false),   // Audio only.
                      AudioVideoTuple(false, true),   // Video only.
                      AudioVideoTuple(true, true)));  // Audio and Video.

// These tests verify that the pipeline handles encrypted media playback
// events (in particular, CDM and license installation) correctly.
class EncryptedAVPipelineImplTest : public ::testing::Test {
 public:
  EncryptedAVPipelineImplTest() {}

  EncryptedAVPipelineImplTest(const EncryptedAVPipelineImplTest&) = delete;
  EncryptedAVPipelineImplTest& operator=(const EncryptedAVPipelineImplTest&) =
      delete;

 protected:
  void SetUp() override {
    pipeline_helper_.reset(new PipelineHelper(true, true, true));
    pipeline_helper_->Setup();
  }

  base::test::TaskEnvironment task_environment_;
  std::unique_ptr<PipelineHelper> pipeline_helper_;
};

// Sets a CDM with license already installed before starting the pipeline.
TEST_F(EncryptedAVPipelineImplTest, SetCdmWithLicenseBeforeStart) {
  base::RepeatingClosure verify_task = base::BindRepeating(
      &VerifyPlay, base::Unretained(pipeline_helper_.get()));
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  pipeline_helper_->SetPipelineStartExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(verify_task)));
  base::RunLoop().Run();
}

// Start the pipeline, then set a CDM with existing license.
TEST_F(EncryptedAVPipelineImplTest, SetCdmWithLicenseAfterStart) {
  base::RepeatingClosure verify_task = base::BindRepeating(
      &VerifyPlay, base::Unretained(pipeline_helper_.get()));
  pipeline_helper_->SetPipelineStartExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(verify_task)));

  task_environment_.RunUntilIdle();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));
  base::RunLoop().Run();
}

// Start the pipeline, set a CDM, and then install the license.
TEST_F(EncryptedAVPipelineImplTest, SetCdmAndInstallLicenseAfterStart) {
  base::RepeatingClosure verify_task = base::BindRepeating(
      &VerifyPlay, base::Unretained(pipeline_helper_.get()));
  pipeline_helper_->SetPipelineStartExpectations();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                std::move(verify_task)));
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));

  task_environment_.RunUntilIdle();
  task_environment_.GetMainThreadTaskRunner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  base::RunLoop().Run();
}

}  // namespace media
}  // namespace chromecast
